## #######################################
## Reanalysis of Network Field Experiment
## #######################################
rm(list=ls())

library(MASS) # ver 7.3-51.5
library(latex2exp)  # ver 0.4.0
library("estimatr") # ver 0.22.0
library("Formula")  # ver 1.2-3

source("estimator.R") # load function
# load data
load(file = "../data/Cui.rdata")

# Only individuals in the second round + NoInfo + have 5 friends
data_r2 <- survey_data_full[survey_data_full$delay == 1 & survey_data_full$info_none == 1, ]
data_r2 <- data_r2[data_r2$num_friend  == 5, ]  # in order to compare  (0.2 vs 0)

# ##################
# ANSE
# ##################
# Compute weights
prob_f <- mean(data_r2$network_rate_preintensive); prob_own <- mean(data_r2$intensive)
prob_indiv <- rep(prob_own, nrow(data_r2))
prob_exp <- c()
for(i in 1:nrow(data_r2)){
  n_t <- data_r2$network_sum_preintensive[i]
  n_c <- data_r2$num_friend[i] - n_t
  prob_exp[i] <- choose(data_r2$num_friend[i], n_t)*(prob_f^n_t)*((1-prob_f)^n_c)
}
prob_use_r2 <- prob_indiv*prob_exp

# Estimation
c_ANSE_model_cov <- ANSE(formula = takeup_survey ~ intensive  | network_rate_preintensive |
                           male + age + agpop + ricearea_2010 + literacy + risk_averse + disaster_prob + 
                           factor(village), 
                         data = data_r2, d = 0, exp = c(0.2, 0),
                         weights = 1/prob_use_r2, type = "model", clusters = data_r2$village)

# Raw estimates 
round(c_ANSE_model_cov$out*100, 2)

# Parametric sensitivity analysis
exp_use <- c(0.2, 0)
coef_U <- c(0.1, 0.2, 0.4)
overlap_use  <- c(0.2, 0.4, 0.6)
para_list <- expand.grid(coef_U, overlap_use)
para_cor <- para_list[,1]*para_list[,2]*(exp_use[1] - exp_use[2])

para_result  <- cbind(c_ANSE_model_cov$out[1] - para_cor, 
                      c_ANSE_model_cov$out[2], 
                      c_ANSE_model_cov$out[1] - para_cor - 1.96*c_ANSE_model_cov$out[2],
                      c_ANSE_model_cov$out[1] - para_cor + 1.96*c_ANSE_model_cov$out[2])

# Nonparametric sensicity analysis
MR_list <- c(1.3, 1.5, 1.8)
RR_list <- c(1.3, 1.5, 1.8)
np_list <- expand.grid(MR_list, RR_list)
lower_mat <-  upper_mat <- matrix(NA, nrow =9,  ncol = 3)
for(j  in 1:9){
  c_ANSE_model_cov_np <- nonpara_ANSE(c_ANSE_model_cov, 
                                      exp_name = "network_rate_preintensive",
                                      exp = exp_use, 
                                      MR_UY = np_list[j,1],  RR_GU = np_list[j,2])
  lower_mat[j,1:3] <- c(c_ANSE_model_cov_np$lower_mean, 
                        c_ANSE_model_cov_np$lower_mean - 1.96*c_ANSE_model_cov_np$lower_se,
                        c_ANSE_model_cov_np$lower_mean + 1.96*c_ANSE_model_cov_np$lower_se)
  upper_mat[j,1:3] <- c(c_ANSE_model_cov_np$upper_mean, 
                        c_ANSE_model_cov_np$upper_mean - 1.96*c_ANSE_model_cov_np$upper_se,
                        c_ANSE_model_cov_np$upper_mean + 1.96*c_ANSE_model_cov_np$upper_se)
}

## Raw Estimate
c_ANSE_model_cov$out
axis_name <- as.character(para_list[,1])
axis_name2 <- as.character(para_list[,2])
ylim_use <- c(-0.15, 0.4)

# All
pdf("../results/Fig_3.pdf", height = 5.5, width  = 12)
par(mfrow = c(1,2), oma = c(0, 1, 0, 0 ))
plot(seq(1:11), c(c_ANSE_model_cov$out[1], NA, para_result[,1]), 
     pch = NA, ylim = ylim_use, 
     col  = c("black", rep("blue", 10)),  xlim = c(0.5, 11.5), 
     xlab =  "", xaxt  = "n", ylab = "",
     main = "Parametric Sensitivity Analysis",  cex.main = 1.55,
     cex.axis  = 1.25)
segments(seq(1:11), c(c_ANSE_model_cov$out[3], NA, para_result[,3]), 
         seq(1:11), c(c_ANSE_model_cov$out[4], NA, para_result[,4]), 
         lwd = 2, col  = c("black", rep("blue", 10)))
points(seq(1:11), c(c_ANSE_model_cov$out[1], NA,  para_result[,1]), 
       pch = 21, bg = "white", cex = 1.5, lwd = 2, 
       col  = c("black", rep("blue", 10)))
abline(h = 0, lty = 2)
text(x = 1, y = 0.18, "Original\nEstimate", font = 1)
mtext(side = 2, text = "Estimated Effects", cex = 1.25, line = 3)
Axis(side = 1, at = seq(from = 3, to = 11), labels = axis_name, cex.axis = 1.0)
Axis(side = 1, at = seq(from = 3, to = 11), labels = axis_name2, line = 1.5, tick = 0, cex.axis = 1.0)
axis(side = 1, at = 2, labels = TeX("$\\lambda$: "), tick = 0, cex.axis = 1.25)
axis(side = 1, at = 2, labels = TeX("$\\pi_{GU}$: "), tick = 0, line = 1.65, cex.axis = 1.25)


## Nonarametric 
axis_name_np <- np_list[,1]
axis_name_np2 <- np_list[,2]
plot(seq(1:11), c(c_ANSE_model_cov$out[1], rep(NA, 10)), 
     pch = NA, ylim = ylim_use, xlim = c(0.5, 11.5), 
     xlab =  "", xaxt  = "n", ylab = "",
     main = "Nonparametric Sensitivity Analysis",  cex.main = 1.55,
     cex.axis  = 1.25)
segments(seq(1), c_ANSE_model_cov$out[3], 
         seq(1), c_ANSE_model_cov$out[4], lwd = 2, col  = "black")
points(seq(1), c_ANSE_model_cov$out[1],
       pch = 21, bg = "white", cex = 1.5, lwd = 2,   col  = c("black"))
arrows(seq(1:11), c(NA, NA, lower_mat[,1]), seq(1:11), c(NA, NA, upper_mat[,1]), 
      lwd = 4, col = "green4", code = 3, length = 0.03, angle = 90)
arrows(seq(1:11), c(NA, NA, lower_mat[,2]), # lower ci of lower bound
       seq(1:11), c(NA, NA, upper_mat[,3]), # upper ci of upper bound
       lwd = 2, col = "green4", code = 3, length = 0.05, angle = 90)
abline(h = 0, lty = 2)
mtext(side = 2, text = "Estimated Effects", cex = 1.25, line = 3)
Axis(side = 1, at = seq(from = 3, to = 11), labels = axis_name_np, cex.axis = 1.0)
Axis(side = 1, at = seq(from = 3, to = 11), labels = axis_name_np2, line = 1.5, tick = 0, cex.axis = 1.0)
Axis(side = 1, at = 2, labels = TeX("$MR_{UY}$: "), tick = 0, cex.axis = 1.0, line =  0.25)
Axis(side = 1, at = 2, labels = TeX("$RR_{GU}$: "), tick = 0, line = 1.65, cex.axis = 1.0)
dev.off()